import torch
from torch import nn
from backbones.text_classifier_model import TextClassifierModel
from dataset.dataloader import generate_vocab

if __name__ == '__main__':
    vocab = generate_vocab()
    vocab_size = len(vocab)
    embedding_dims = 256
    device = torch.device("cuda")
    model = TextClassifierModel(vocab_size, embedding_dims, 8, 2)
    model = model.to(device)
    model.load_state_dict(torch.load("./save/best.pt", weights_only=False))

    model.eval()

    sentence = "不太满意，单鞋不值这个价"
    sentence_idx = [vocab.to_idx('<CLS>')] + [vocab.to_idx(token) for token in list(sentence)] + [vocab.to_idx('<SEP>')]
    inputs = torch.tensor([sentence_idx], device=device)

    outputs = model(inputs, None)
    print(outputs)